# -*- coding:utf-8 -*-
#common libs

import numpy as np

def traceSeries(array):
    '''
        generate trace series along the third axis
    '''
    assert(len(array.shape) ==3)
    trs = np.zeros((array.shape[2],))
    for i in range(0,array.shape[2]):
        trs[i]    = np.trace(array[:,:,i])
    return trs
def direction(p1,p2):
    '''
    direction of nodes p1 and p2. Support 2D and 3D, return the angle from p2 to p1
    In 2D case, theta will be a scaler indicates the angle from p2 to p1, which ranges from -pi~pi;
    In 3D case theta will be a vector, theta(1) is azimuthal angle (-pi~ pi)  and theta(2) is altitude angle(-pi/2~pi/2)
    '''
    if p1.size ==2:
        v = p1 -p2 
        theta =  np.arctan2(v[1],v[0])
    else:
        v = p1 -p2 
        origin = np.zeros((2,))
        vxy = v[0:2]
        theta = direction(vxy,origin)
        r  =  np.sqrt(v[0]**2+v[1]**2)
        vrz =  np.array([r,v[2]],dtype=float)
        phi = direction(vrz,origin)
        theta = [theta,phi]
    return theta
def angletransform(theta,phi):
    '''
        thetaprime is angle around z axis
    '''
    # thetaprime = np.arccos(np.cos(phi)* np.cos(theta))
    # phiprime   = np.arccos((np.cos(phi)* np.sin(theta))/(np.sqrt((np.cos(phi)* np.sin(theta))**2 + np.sin(phi)**2)))
    thetaprime = np.arctan2(np.sqrt(np.sin(phi)**2+(np.cos(phi)*np.sin(theta))**2),np.cos(phi)*np.cos(theta))
    phiprime   = np.arctan2(np.sin(phi),np.cos(phi)*np.sin(theta))
    return [thetaprime,phiprime]
def rotationMatrix(theta,phi=None,inverse=None):
    '''
        X  = R X'
    '''
    if phi is None:
        dim2 = True
    else:
        dim2 = False

    if inverse is None:
        inverse  = False
    
    if dim2 == False:
        angleprime = angletransform(theta,phi)
        t   = angleprime[0]
        p   = angleprime[1]
        
        if inverse==False:
            Rx = np.array(((1, 0, 0), (0, np.cos(p), -np.sin(p)), (0, np.sin(p), np.cos(p))),dtype=float)
            Rz = np.array((( np.cos(t), -np.sin(t),0), (np.sin(t), np.cos(t),0),( 0, 0,1)),dtype=float)
            rotationMatrix = np.dot(Rx,Rz)
        else:
            Rx = np.array(((1, 0, 0), (0, np.cos(-p), -np.sin(-p)), (0, np.sin(-p), np.cos(-p))),dtype=float)
            Rz = np.array((( np.cos(-t), -np.sin(-t),0), (np.sin(-t), np.cos(-t),0),( 0, 0,1)),dtype=float)
            rotationMatrix = np.dot(Rz,Rx)
    else:
        if inverse==True:
            theta = - theta 
        rotationMatrix = np.array(((np.cos(theta),-np.sin(theta)),(np.sin(theta),np.cos(theta))))
    return rotationMatrix

if __name__ == '__main__':
    a = np.random.randn(2,2,3)
    # print a
    # print traceSeries(a)
    print direction(np.array((1,1,1)),np.array((0,0,0)))
    # print angletransform(np.pi/4,np.arccos(np.sqrt(2.0/3)))
    # print rotationMatrix(np.pi/4,np.arccos(np.sqrt(2.0/3)))
    # print np.dot(rotationMatrix(np.pi/4,np.arccos(np.sqrt(2.0/3))),np.array((np.sqrt(3),0,0)))